# python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# References:
# DeiT: https://github.com/facebookresearch/deit
# BEiT: https://github.com/microsoft/unilm/tree/master/beit
# --------------------------------------------------------
import argparse  # 导入argparse库，用于解析命令行参数
import datetime  # 导入datetime库，用于处理日期和时间
import json  # 导入json库，用于处理JSON数据
import numpy as np  # 导入numpy库，用于科学计算
import os  # 导入os库，用于与操作系统交互
import time  # 导入time库，用于处理时间相关功能
from pathlib import Path  # 从pathlib模块导入Path类，用于处理文件系统路径
import torch.nn as nn  # 导入torch.nn模块，用于构建神经网络

import torch  # 导入torch库，PyTorch的核心库
import torch.backends.cudnn as cudnn  # 导入cudnn模块，用于配置CUDA深度神经网络库
from torch.utils.tensorboard import SummaryWriter  # 从torch.utils.tensorboard导入SummaryWriter类，用于记录TensorBoard日志
import torchvision.transforms as transforms  # 导入torchvision.transforms模块，用于图像预处理
import torchvision.datasets as datasets  # 导入torchvision.datasets模块，用于加载常用数据集

import timm  # 导入timm库，用于加载预训练模型和优化器

# assert timm.__version__ == "0.3.2"  # 检查timm库的版本，确保为0.3.2（此行代码被注释掉了）
import timm.optim.optim_factory as optim_factory  # 导入timm的优化器工厂模块

import util.misc as misc  # 导入自定义的util.misc模块，包含一些实用函数
from util.misc import NativeScalerWithGradNormCount as NativeScaler  # 从util.misc导入NativeScalerWithGradNormCount类，并重命名为NativeScaler

import models_Mambamae  # 导入自定义的models_Mambamae模块，包含MedMamba模型的定义
from VMamba.classification.medical.medmamba import VSSM  # 从VMamba.classification.medical.medmamba导入VSSM类，用于加载Swin-UMamba模型
#from VMamba.classification.models.vmamba import VSSM

from engine_pretrain_vmamba import train_one_epoch_MedMamba  # 从engine_pretrain模块导入train_one_epoch_MedMamba函数，用于MedMamba模型的预训练


def get_args_parser():
    """定义命令行参数解析器"""
    parser = argparse.ArgumentParser('medmambaMAE pre-training', add_help=False)
    # 添加命令行参数，包括批次大小、训练轮数、梯度累积迭代次数等
    parser.add_argument('--batch_size', default=64, type=int,
                        help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus') #64
    parser.add_argument('--epochs', default=300, type=int)
    parser.add_argument('--accum_iter', default=1, type=int,
                        help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')

    # 添加模型参数，包括模型名称、输入图像大小、掩码比率等
    parser.add_argument('--model', default='mae_vit_base_patch16', type=str, metavar='MODEL',
                        help='Name of model to train')
    parser.add_argument('--input_size', default=224, type=int,
                        help='images input size')
    parser.add_argument('--mask_ratio', default=0.75, type=float,
                        help='Masking ratio (percentage of removed patches).')
    parser.add_argument('--norm_pix_loss', action='store_true',
                        help='Use (per-patch) normalized pixels as targets for computing loss')
    parser.set_defaults(norm_pix_loss=False)

    # 添加优化器参数，包括权重衰减、学习率等
    parser.add_argument('--weight_decay', type=float, default=0.05,
                        help='weight decay (default: 0.05)')
    parser.add_argument('--lr', type=float, default=None, metavar='LR',
                        help='learning rate (absolute lr)')
    parser.add_argument('--blr', type=float, default=1.5e-4, metavar='LR',
                        help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')  #1.5e-4
    parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
                        help='lower lr bound for cyclic schedulers that hit 0')
    parser.add_argument('--warmup_epochs', type=int, default=20, metavar='N',
                        help='epochs to warmup LR') #40

    # 添加数据集参数，包括数据集路径、输出目录、日志目录等
    parser.add_argument('--data_path', default='TN3K_initial/class', type=str,
                        help='dataset path')
    parser.add_argument('--output_dir', default='./output/MedMamba/pretrain/TN3K_initial/bs_64_blr_1.5e-4_wd_0.05_select_layer_cv_newtransform_75%',
                        help='path where to save, empty for no saving')
    parser.add_argument('--log_dir', default='./output/MedMamba/pretrain/TN3K_initial/bs_64_blr_1.5e-4_wd_0.05_select_layer_cv_newtransform_75%',
                        help='path where to tensorboard log')
    parser.add_argument('--device', default='cuda',
                        help='device to use for training / testing')
    parser.add_argument('--seed', default=0, type=int)
    parser.add_argument('--resume', default='',
                        help='resume from checkpoint')

    # 添加训练起始轮数、工作线程数、内存固定选项等参数
    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                        help='start epoch')
    parser.add_argument('--num_workers', default=8, type=int)
    parser.add_argument('--pin_mem', action='store_true',
                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
    parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem')
    parser.set_defaults(pin_mem=True)

    # 添加分布式训练参数，包括世界大小、本地排名、分布式训练URL等
    parser.add_argument('--world_size', default=1, type=int,
                        help='number of distributed processes')
    parser.add_argument('--local_rank', default=-1, type=int)
    parser.add_argument('--dist_on_itp', action='store_true')
    parser.add_argument('--dist_url', default='env://',
                        help='url used to set up distributed training')

    return parser


def load_medmamba_model(device='cuda'):
    """加载Swin-UMamba模型"""
    # 创建VSSM类的实例，初始化模型结构
    model = VSSM(depths=[2, 2, 4, 2], dims=[96, 192, 384, 768], num_classes=4).to("cuda")
    model = model.cuda()
    # 指定预训练权重文件的路径
    checkpoint_path = "/home/ma-user/work/mae/check_point/MedMamba.pth"
    # 加载预训练权重
    checkpoint = torch.load(checkpoint_path, map_location=device)
    # 将预训练权重加载到模型中
    model.load_state_dict(checkpoint, strict=False)  # 根据需要选择是否要 strict

    return model


def main(args):
    """主函数，执行MedMamba模型的预训练"""
    # 初始化分布式训练模式
    misc.init_distributed_mode(args)  # 初始化分布式训练模式

    # 打印当前工作目录
    print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))  # 打印当前工作目录
    # 打印命令行参数
    print("{}".format(args).replace(', ', ',\n'))  # 打印命令行参数

    # 设置设备（CPU或GPU）
    device = torch.device(args.device)  # 设置设备（CPU或GPU）

    # 设置随机种子以保证结果的可复现性
    seed = args.seed + misc.get_rank()
    torch.manual_seed(seed)
    np.random.seed(seed)

    # 开启cudnn的benchmark模式以加速计算
    cudnn.benchmark = True  # 开启cudnn的benchmark模式以加速计算

    # 定义图像预处理流程
    transform_train = transforms.Compose([
        # 对图像进行随机裁剪和缩放
        transforms.RandomResizedCrop(args.input_size, scale=(0.5, 1.0), interpolation=3),
        # 对图像进行随机水平翻转
        transforms.RandomHorizontalFlip(),
        # 将图像转换为张量
        transforms.ToTensor(),
        # 对图像进行标准化处理
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
    # 加载训练数据集
    dataset_train = datasets.ImageFolder(os.path.join(args.data_path, ''), transform=transform_train)
    #print(dataset_train, dataset_train.shape)  # 打印数据集信息

    # 根据是否进行分布式训练，设置不同的采样器
    if True:  # args.distributed:
        num_tasks = misc.get_world_size()
        global_rank = misc.get_rank()
        sampler_train = torch.utils.data.DistributedSampler(
            dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
        )
        print("Sampler_train = %s" % str(sampler_train))
    else:
        sampler_train = torch.utils.data.RandomSampler(dataset_train)

    # 如果全局排名为0且指定了日志目录，则创建日志目录并初始化TensorBoard日志记录器
    if global_rank == 0 and args.log_dir is not None:
        os.makedirs(args.log_dir, exist_ok=True)
        log_writer = SummaryWriter(log_dir=args.log_dir)
    else:
        log_writer = None

    # 创建数据加载器
    data_loader_train = torch.utils.data.DataLoader(
        dataset_train, sampler=sampler_train,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        pin_memory=args.pin_mem,
        drop_last=True,
    )
    
    # 加载Swin-UMamba模型
    medmamba = load_medmamba_model()
    medmamba.to(device)
    medmamba.eval()  # 将模型设置为评估模式
    
    # 定义要训练的模型
    model = models_Mambamae.__dict__[args.model](norm_pix_loss=args.norm_pix_loss)

    model.to(device)  # 将模型移动到指定设备

    model_without_ddp = model  # 保存没有DDP包装的模型引用

    # 计算有效批次大小
    eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
    
    # 如果只指定了基础学习率，则根据有效批次大小计算实际学习率
    if args.lr is None:  # only base_lr is specified
        args.lr = args.blr * eff_batch_size / 256

    # 打印基础学习率和实际学习率
    print("base lr: %.2e" % (args.lr * 256 / eff_batch_size))
    print("actual lr: %.2e" % args.lr)

    # 打印梯度累积迭代次数和有效批次大小
    print("accumulate grad iterations: %d" % args.accum_iter)
    print("effective batch size: %d" % eff_batch_size)

    # 如果进行分布式训练，则使用DistributedDataParallel包装模型
    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
        model_without_ddp = model.module  # 获取DDP包装内部的模型引用
    
    # 设置优化器，并为偏置和归一化层设置权重衰减为0
    param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
    optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
    print(optimizer)  # 打印优化器信息
    loss_scaler = NativeScaler()  # 初始化损失缩放器

    # 加载模型权重、优化器状态和损失缩放器状态（如果指定了恢复路径）
    misc.load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler)

    print(f"Start training for {args.epochs} epochs")  # 打印训练开始信息
    start_time = time.time()  # 记录训练开始时间
    for epoch in range(args.start_epoch, args.epochs):  # 循环训练指定轮数
        if args.distributed:
            # 设置采样器的当前轮数（分布式训练时使用）
            data_loader_train.sampler.set_epoch(epoch)  # 设置采样器的当前轮数（分布式训练时使用）
        # 执行一轮训练
        train_stats = train_one_epoch_MedMamba(model, medmamba, data_loader_train,
                                        optimizer, device, epoch, loss_scaler,
                                        log_writer=log_writer, args=args)
        # 在训练结束时保存模型（如果指定了输出目录且达到最后一轮）
        if args.output_dir and epoch + 1 == args.epochs:
            misc.save_model(
                args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
                loss_scaler=loss_scaler, epoch=epoch)

        # 记录训练日志
        log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
                        'epoch': epoch,}

        if args.output_dir and misc.is_main_process():
            if log_writer is not None:
                log_writer.flush()  # 刷新TensorBoard日志记录器
            with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f:
                f.write(json.dumps(log_stats) + "\n")  # 将日志写入文件

    total_time = time.time() - start_time  # 计算总训练时间
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))  # 将总训练时间格式化为字符串
    print('Training time {}'.format(total_time_str))


if __name__ == '__main__':
    # 解析命令行参数
    args = get_args_parser()
    args = args.parse_args()
    # 如果指定了输出目录，则创建该目录
    if args.output_dir:
        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    # 执行主函数开始训练
    main(args)
